#!/usr/bin/env python

import sys
sys.path.append('..')
import numpy as np
from datetime import datetime, timedelta
from optparse import OptionParser

from INFILE import maindir, siteid, webdir, wfoid
sys.path.append('%s/data_parsers' % maindir)


# Use the command line to get the variable to compute, the model and the station ID
parser = OptionParser()
parser.add_option('-s', '--site', dest='siteid', action='store', type='string', default=siteid.upper(), help='4-letter id of desired site')
parser.add_option('-m', '--model', dest='model', action='store', type='string', default='GFS', help='Model to use (GFS, NAM, NWS)')
parser.add_option('-v', '--var', dest='variable', action='store', type='string', default='high', help='Variable to compute (high, low, wind)')
parser.add_option('-e', '--export', dest='export', action='store_true', default=False, help='Flag to write images to web or not')

(opts,args) = parser.parse_args()

invar = opts.variable
model_name = opts.model.upper()
station = opts.siteid.upper()
export_flag = opts.export

if model_name in ['USL12Z','USL22Z','NWS']:
    length = [10]
else:
    length = [10,30,60]

use_KF = True
use_GFSE = False

import matplotlib
if export_flag:
    matplotlib.use('agg')
    import os

def main():
    # First, make sure the truth is updated to the most recent dates
    from surface_parse_bufkit import get_todays_verification
    from encode_truth import update_truth
    verify_dict = get_todays_verification(siteid,wfoid)
    daykey = verify_dict.keys()[0]
    daynum = daykey.day
    high = float('%2.0f' % verify_dict[daykey][0])
    low = float('%2.0f' % verify_dict[daykey][1])
    wind = float('%2.0f' % verify_dict[daykey][3])
    precip = float('%1.2f' % float(verify_dict[daykey][4]))
    print "Updating today:", high, low, wind, precip
    update_truth(station,wfoid,daynum,high,low,wind,precip)
    truth, model = get_verification(model_name,station,length)

    # Import the Kalman Filter adjustment if desired
    if use_KF:
        from kf_adjustment import run_filter

    # Import the GFSE values if requested
    if use_GFSE:
        from get_gfse import get_gfse

    # If all variables requested (default) loop through each
    if invar == 'all':
        # Plot all three major variables
        for var in ['high','low','wind']:
            if use_KF:
                print "Computing KF values for %s" % var
                kfdates,kfvals = run_filter(truth,model,max(length),var)
                kf_d = {}
                for d,v in zip(kfdates,kfvals):
                    kf_d[d] = v
            else:
                kf_d = None
            if use_GFSE and var in ['high','low']:
                gfsed = {}
                gfsedates,bma,low50,high50,low90,high90 = get_gfse(station,var+'s')
                for d,b,l5,h5,l9,h9 in zip(gfsedates,bma,low50,high50,low90,high90):
                    gfsed[d] = {}
                    gfsed[d]['bma'] = b
                    gfsed[d]['low50'] = l5
                    gfsed[d]['high50'] = h5
                    gfsed[d]['low90'] = l9
                    gfsed[d]['high90'] = h9
            else:
                gfsed = None

            plot_verification(truth,model,var,model_name,station,length,kf_vals=kf_d,gfse_vals=gfsed)

    else:
        if use_KF:
            print "Computing KF values for %s" % invar
            kfdates,kfvals = run_filter(truth,model,max(length),invar)
            kf_d = {}
            for d,v in zip(kfdates,kfvals):
                kf_d[d] = v
        else:
            kf_d = None

        plot_verification(truth,model,invar,model_name,station,length, kf_vals=kf_d)


def plot_verification(actual,fcst,var,mod_name,site,length,kf_vals=None,gfse_vals=None):
    """ Make the verification plots """
    import matplotlib.pyplot as plt
    # Loop through each length desired
    for flen in length:
        # Build the date lists
        model_dates = fcst.keys()
        model_dates.sort()
        model_dates = model_dates[-flen:]

        truth_dates = actual.keys()
        truth_dates.sort()
        truth_dates = truth_dates[truth_dates.index(model_dates[0]):]

        # Now build the plotting list for each variable
        model = [fcst[d][var] for d in model_dates]
        truth = [actual[d][var] for d in truth_dates]

        # Start making the figure
        plt.figure(figsize=(12,11))
        plt.subplot(2,1,1)

        # Check if we have GFSE background to plot
        if gfse_vals != None:
            # Make list of plotting vals
            gfse_dates = gfse_vals.keys()
            bma_list = [gfse_vals[d]['bma'] for d in gfse_dates]
            low_50 = [gfse_vals[d]['low50'] for d in gfse_dates]
            high_50 = [gfse_vals[d]['high50'] for d in gfse_dates]
            low_90 = [gfse_vals[d]['low90'] for d in gfse_dates]
            high_90 = [gfse_vals[d]['high90'] for d in gfse_dates]
            # Plot between the intervals and plot the bma
            plt.plot(gfse_dates,bma_list, 'g-', linewidth=2, linestyle='--', label='GFSE BMA')
            plt.fill_between(gfse_dates,high_90,low_90,color='lightgreen', alpha=0.2, linewidth=1)
            plt.fill_between(gfse_dates,high_50,low_50,color='lightgreen', alpha=0.5, linewidth=1)



        # Check to see if we have the Kalman filter values to plot
        if kf_vals != None:
            kf_dates = kf_vals.keys()
            kf_dates.sort()
            # Match this to the actual plotting length
            kf_dates = kf_dates[kf_dates.index(model_dates[0]):]
            kf_fcst = [kf_vals[d] for d in kf_dates]
            # Plot the KF correction
            plt.plot(kf_dates, kf_fcst, 'm-', linewidth=2, linestyle='--', label= 'KF')
            tom_kf = kf_fcst[-1]
        else:
            tom_kf = None

        # Plot the actual model and truth values
        plt.plot(model_dates, model, 'b-', linewidth=3, label=mod_name)
        plt.plot(truth_dates, truth, 'k-', linewidth=4, label='Observ.')

        # Make annotations for tomorrow's model forecast
        tom_mod = model[-1]
        # Also do the annotation for the KF if it exists
        if tom_kf != None:
            # Figure out which one is higher
            if tom_mod >= tom_kf:
                mosy = 25
            else:
                mosy = -25

            if var == 'winds':
                mosy = abs(mosy)
                KFy = mosy
            else:
                KFy = -1 * mosy

            if var == 'precip':
                tom_kf_str = '%3.2f' % tom_kf
            else:
                tom_kf_str = '%3.1f' % tom_kf
            plt.annotate(tom_kf_str, xy=(matplotlib.dates.date2num(model_dates[-1]),tom_kf), xycoords='data',\
                textcoords='offset points', xytext=(-10,KFy), size='medium', va='center', ha='left',\
                bbox=dict(boxstyle='round4', fc='w', edgecolor='m'), arrowprops=dict(arrowstyle='->', edgecolor='m'))
        else:
            mosy = 25

        # If the variable is precip, then go out to two decimal places.  Otherwise, just go to one
        if var == 'precip':
            tom_str = '%3.2f' % tom_mod
        else:
            tom_str = '%3.1f' % tom_mod
        plt.annotate(tom_str, xy=(matplotlib.dates.date2num(model_dates[-1]),tom_mod), xycoords='data',\
            textcoords='offset points', xytext=(-10,mosy), size='medium', va='center', ha='left',\
            bbox=dict(boxstyle='round4', fc='w', edgecolor='b'), arrowprops=dict(arrowstyle='->', edgecolor='b'))

        # Now make the legend
        leg = plt.legend(loc=0,ncol=2)
        ltexts = leg.get_texts()
        plt.setp(ltexts,fontsize='small')
        
        # Write the title
        plt.title('Past %d-day forecast %s comparison at %s (12Z %s)' % (flen, var, site.upper(), mod_name))

        # Now some bookkeeping adjustments
        # If plotting the wind or precip, start the y-axis at zero (no negatives)
        if var in ['wind','precip']:
            ymin,ymax = plt.ylim()
            if ymin < 0.0:
                plt.ylim(ymin=0.0)

        # Format the axes
        # Make the dates look nice
        ax = plt.gca()
        ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%b %d'))
        plt.grid()
        # Set the correct axis label
        if var == 'wind':
            plt.ylabel('Wind speed (kts)')
        elif var == 'precip':
            plt.ylabel('Precipitation (in.)')
        else:
            plt.ylabel('Temperature (F)')

        # That wraps up panel one.  Now on to the bar chart for panel two
        plt.subplot(2,1,2)
        # Simply find the model error by subtracting from truth
        # Make sure that the truth and model values line up
        model_verified = []
        truth_verified = []
        bar_dates = []
        for t,tdate in zip(truth, truth_dates):
            if tdate in model_dates:
                truth_verified.append(t)
                model_verified.append(model[model_dates.index(tdate)])
                bar_dates.append(tdate)
        model_error = np.subtract(model_verified,truth_verified)
        model_error = list(model_error)
        # Append a nans to the list to make it line up with the current date
        barend = bar_dates[-1]
        while barend < model_dates[-1]:
            barend = barend + timedelta(days=1)
            bar_dates.append(barend)
            model_error.append(np.nan)
        model_error = np.array(model_error)
        # Compute the mean error
        masked_model = np.ma.masked_array(model_error,np.isnan(model_error))
        meanf_model_error = np.mean(masked_model)
        # Format this with a + or - sign
        if meanf_model_error >= 0.0:
            mean_model_error = '+%2.1f' % meanf_model_error
        else:
            mean_model_error = '%2.1f' % meanf_model_error

        # Start building text for output
        txtb = 'Mean %s error: %s' % (mod_name, mean_model_error)

        # Now make the dates work for the bar chart
        dxval = abs(matplotlib.dates.date2num(bar_dates[-1]) - matplotlib.dates.date2num(bar_dates[-2]))
        bar_dates_adj = [datetime(d.year,d.month,d.day,0) - timedelta(hours=12) for d in bar_dates]

        # Plot the bar chart
        model_error_bar = plt.bar(bar_dates_adj, model_error, color='b', align='edge', bottom = 0, width=dxval, label=mod_name)
        
        if gfse_vals != None:
            pass

        # And check for Kalman filter values
        if kf_vals != None:
            kf_verified = [kf_vals[d] if d in kf_vals.keys() else np.nan for d in bar_dates]
            kf_verified = kf_verified[:-1]
            kf_error = np.subtract(kf_verified, truth_verified)
           
            # Find the masked mean
            masked_kf = np.ma.masked_array(kf_error, np.isnan(kf_error))
            meanf_kf_error = np.mean(masked_kf)
            # Format the error with + or -
            if meanf_kf_error >= 0.0:
                mean_kf_error = '+%2.1f' % meanf_kf_error
            else:
                mean_kf_error = '%2.1f' % meanf_kf_error
            txtb = ' | '.join([txtb,'Mean KF error: %s' % mean_kf_error])

            # Now plot the bar chart
            bar_dates_kf = [datetime(d.year,d.month,d.day,0) - timedelta(hours=6) for d in bar_dates]
            bar_dates_kf = bar_dates_kf[:-1]
            kf_error_bar = plt.bar(bar_dates_kf, kf_error, color='m', width=dxval/2., alpha=0.3, label='KF error')


    
        # Fix the plot by setting the beginning and ending points
        plt.xlim(xmax=matplotlib.dates.date2num(bar_dates_adj[-1]+timedelta(hours=12)))
        plt.xlim(xmin=matplotlib.dates.date2num(bar_dates_adj[0]+timedelta(hours=12)))
        
        # Adjust the buffer around the bars
        yn,yx = plt.ylim()
        if var not in ['wind']:
            plt.ylim((yn-2,yx+1))
        else:
            plt.ylim((yn-2,yx+7))

        # Make a zero line
        zero_line = plt.axhline(y=0, color='k', linestyle='--', linewidth=4, label='Zero line')

        # Make the text box with the mean errors
        plt.text(0.18, -1.15, txtb, fontsize=14, transform=ax.transAxes, bbox=dict(facecolor='gray',alpha=0.1))

        # Make the legend and title
        leg = plt.legend(loc=0)
        ltexts = leg.get_texts()
        plt.setp(ltexts,fontsize='small')

        # Set the locators
        ax = plt.gca()
        ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%b %d'))
        plt.grid()
        plt.ylabel('%s error' % var)
        
        # And that's it for now
        if export_flag:
            print "Exporting %s %s %d day verif." % (mod_name, var, flen)
            plt.savefig('%s_%s_%s_%dday.png' % (site, mod_name, var, flen), bbox_inches='tight')
            os.system('mv %s_%s_%s_%dday.png %s' % (site, mod_name, var, flen, webdir))
        else:
            plt.show()
        plt.close()


def get_verification(model,site,length):
    """ Get the verification time series (both the truth and the model forecast)
    for the specified site, variable, and verification length """
    # If we have the GFS or NAM, read from their archive
    if model_name.startswith('GFS') or model_name.startswith('NAM'):
        from mos_archive_reader import read_mos_dict, mosval
        modeld = read_mos_dict('%s/site_data/%s_%s_MOS.pickle' % (maindir,site.upper(),model))
    else:
        # These are already in dictionary form and don't need a special import
        import cPickle
        infile = open('%s/site_data/%s_%s_fcst.pickle' % (maindir,site.upper(),model),'r')
        modeld = cPickle.load(infile)

    # Model forecast dates
    dates = modeld.keys()
    # Be sure these dates are set to 0Z
    newdates = [d.date() for d in dates]
    dates = newdates
    dates.sort()

    # Now get the truth values
    from encode_truth import encode_truth
    actual_values = encode_truth(site.upper())

    # Now return only the most recent (max length) dates
    good_dates = dates[-max(length):]


    # Delete all values not in good_dates
    for d in actual_values.keys():
        if d.date() not in good_dates:
            actual_values.pop(d)
        else:
            actual_values[d.date()] = actual_values[d]
            actual_values.pop(d)
    for d in modeld.keys():
        if d.date() not in good_dates:
            modeld.pop(d)
        else:
            modeld[d.date()] = modeld[d]
            modeld.pop(d)
    # Return the two dictionaries
    return actual_values, modeld



if __name__ == '__main__':
    main()
